# -*- coding: utf-8 -*-

import numpy as np
from torchvision import transforms
import random
import torch


def random_transform(x):
    trans_args = {}
    
    if np.random.randint(0, 1) == 0:
        x = transforms.functional.hflip(x)
        trans_args['hflip'] = True
    else:
        trans_args['hflip'] = False
    
    # if np.random.randint(0, 1) == 0:
    #     x = transforms.functional.vflip(x)
        
    scale = np.random.uniform(0.8, 1.2)
    size = int(scale * 512)
    x = transforms.functional.resize(x, size)
    if size > 512:
        x = transforms.functional.center_crop(x, 512)
    else:
        left = int((512 - size) / 2)
        right = (512 - size) - left
        x = transforms.functional.pad(x, [left, left, right, right], 170)
    trans_args['scale'] = scale
        
    rot = random.randint(0, 3)
    x = torch.rot90(x, rot, dims=(-2, -1))
    trans_args['rot'] = rot
    
    return x, trans_args


def target_transform(x, trans_args):
    if trans_args['hflip']:
        x = transforms.functional.hflip(x)
        
    scale = trans_args['scale']
    size = int(scale * 512)
    x = transforms.functional.resize(x, size)
    if size > 512:
        x = transforms.functional.center_crop(x, 512)
    else:
        left = int((512 - size) / 2)
        right = (512 - size) - left
        x = transforms.functional.pad(x, [left, left, right, right], 0)
        
    rot = trans_args['rot']
    x = torch.rot90(x, rot, dims=(-2, -1))
    
    return x



# if __name__ == "__main__":
#     from CC_CCI_dataset import LabeledDataset, visualize
#     # import matplotlib.pyplot as plt
#     labeled_dataset = LabeledDataset("D:\data\ct_lesion_seg\image", "D:\data\ct_lesion_seg\mask")
#     x, y = labeled_dataset[113]
#     print(x.shape)
#     print(y.shape)
    
#     visualize(x, y)
#     x = torch.unsqueeze(x, dim=0)
#     y = torch.unsqueeze(y, dim=0)
#     t, args = random_transform(x)
#     s = target_transform(y, args)
#     print(args)